'''
Structure generation by evolutionary algorithm
'''

from ..struc_util import out_poscar


class EA_generation:
    '''
    generate structures by evolutionary algorithm

    # ---------- args
    sp: instance of Select_parents class
    symprec (float): precision for symmetry finding
    id_start (int or None): starting id
                            if None, id_start = max(self.fitness.keys()) + 1
    init_pos_path (str or None): if not None, structure data in POSCAR format
                                     is appended to init_pos_path

    # ---------- instance methods
    self.gen_crossover(self, n_crsov, co)

    self.gen_permutation(self, n_perm, pm)

    self.gen_strain(self, n_strain, st)
    '''

    def __init__(self, sp, symprec=0.001, id_start=None, init_pos_path=None):
        # ---------- check args
        # ------ sp
        if (hasattr(sp, 'cum_fit') or hasattr(sp, 't_size')):
            self.sp = sp
        else:
            raise ValueError('sp is wrong')
        # ------ symprec
        if isinstance(symprec, float):
            if symprec < 0.0:
                raise ValueError('symprec must not be negative value')
            else:
                self.symprec = symprec
        else:
            raise TypeError('symprec must be float')
        # ------ id_offset
        if id_start is None:
            self.cid = max(sp.fitness.keys()) + 1
        elif isinstance(id_start, int):
            if id_start < (max(sp.fitness.keys()) + 1):
                raise ValueError('id_start is already included'
                                 ' structure ID of the data')
            else:
                self.cid = id_start
        else:
            raise TypeError('id_start must be int or None')
        # ------ init_pos_path
        if init_pos_path is not None:
            if isinstance(init_pos_path, str):
                self.init_pos_path = init_pos_path
            else:
                raise TypeError('init_pos_path must be str or None')
        # ---------- initialize data
        self.offspring = {}    # structure data
        self.parents = {}    # tuple of parents ID
        self.operation = {}

    def gen_crossover(self, n_crsov, co):
        '''
        generate structures by crossover

        # ---------- args
        n_crsov (int): number of structures generated by crossover
        co: instance of Crossover class
        '''
        # ---------- check args
        # ------ n_crsov
        if isinstance(n_crsov, int):
            if n_crsov <= 0:
                raise ValueError('n_crsov must be positive int')
        else:
            raise TypeError('n_crsov must be int')
        # ---------- generate structures by crossover
        struc_cnt = 0
        while struc_cnt < n_crsov:
            # ------ select parents
            pid_A, pid_B = self.sp.get_parents(n_parent=2)
            # ------ generate child
            child = co.gen_child(self.sp.struc_data[pid_A],
                                 self.sp.struc_data[pid_B])
            # ------ success
            if child is not None:
                self.offspring[self.cid] = child
                self.parents[self.cid] = (pid_A, pid_B)
                self.operation[self.cid] = 'crossover'
                try:
                    spg_sym, spg_num = child.get_space_group_info(
                        symprec=self.symprec)
                except TypeError:
                    spg_num = 0
                    spg_sym = None
                print('Structure ID {0:>6} was generated'
                      ' from {1:>6} and {2:>6} by crossover.'
                      ' Space group: {3:>3} {4}'.format(self.cid, pid_A, pid_B,
                                                        spg_num, spg_sym))
                if self.init_pos_path is not None:
                    out_poscar(child, self.cid, self.init_pos_path)
                self.cid += 1
                struc_cnt += 1

    def gen_permutation(self, n_perm, pm):
        '''
        generate structures by permutation

        # ---------- args
        n_perm (int): number of structures generated by permutation
        pm: instance of Permutation class
        '''
        # ---------- check args
        # ------ n_perm
        if isinstance(n_perm, int):
            if n_perm <= 0:
                raise ValueError('n_perm must be positive int')
        else:
            raise TypeError('n_perm must be int')
        # ------ pm
        if not hasattr(pm, 'ntimes'):
            raise ValueError('pm is wrong')
        # ---------- generate structures by permutation
        struc_cnt = 0
        while struc_cnt < n_perm:
            # ------ select parents
            pid, = self.sp.get_parents(n_parent=1)    # comma for list[0]
            # ------ generate child
            child = pm.gen_child(self.sp.struc_data[pid])
            # ------ success
            if child is not None:
                self.offspring[self.cid] = child
                self.parents[self.cid] = (pid, )    # tuple
                self.operation[self.cid] = 'permutation'
                try:
                    spg_sym, spg_num = child.get_space_group_info(
                        symprec=self.symprec)
                except TypeError:
                    spg_num = 0
                    spg_sym = None
                print('Structure ID {0:>6} was generated'
                      ' from {1:>6} by permutation.'
                      ' Space group: {2:>3} {3}'.format(self.cid, pid,
                                                        spg_num, spg_sym))
                if self.init_pos_path is not None:
                    out_poscar(child, self.cid, self.init_pos_path)
                self.cid += 1
                struc_cnt += 1

    def gen_strain(self, n_strain, st):
        '''
        generate structures by strain

        # ---------- args
        n_strain (int): number of structures generated by strain
        st: instance of Strain class
        '''
        # ---------- check args
        # ------ n_strain
        if isinstance(n_strain, int):
            if n_strain <= 0:
                raise ValueError('n_strain must be positive int')
        else:
            raise TypeError('n_strain must be int')
        # ------ st
        if not hasattr(st, 'sigma'):
            raise ValueError('st is wrong')
        # ---------- generate structures by strain
        struc_cnt = 0
        while struc_cnt < n_strain:
            # ------ select parents
            pid, = self.sp.get_parents(n_parent=1)    # comma for list[0]
            # ------ generate child
            child = st.gen_child(self.sp.struc_data[pid])
            # ------ success
            if child is not None:
                self.offspring[self.cid] = child
                self.parents[self.cid] = (pid, )    # tuple
                self.operation[self.cid] = 'strain'
                try:
                    spg_sym, spg_num = child.get_space_group_info(
                        symprec=self.symprec)
                except TypeError:
                    spg_num = 0
                    spg_sym = None
                print('Structure ID {0:>6} was generated'
                      ' from {1:>6} by strain.'
                      ' Space group: {2:>3} {3}'.format(self.cid, pid,
                                                        spg_num, spg_sym))
                if self.init_pos_path is not None:
                    out_poscar(child, self.cid, self.init_pos_path)
                self.cid += 1
                struc_cnt += 1
